import json
from models.model import OpinionModel
import matplotlib.pyplot as plt
import numpy as np

def run_simulation(num_steps: int, num_agents: int, num_groups: int, network_density: float):
    """运行模拟"""
    model = OpinionModel(num_agents, num_groups, network_density)
    
    # 记录模拟结果
    results = {
        "opinions": [],
        "group_opinions": [[] for _ in range(num_groups)]
    }
    
    for step in range(num_steps):
        state = model.step()
        
        # 记录所有智能体的观点
        opinions = [agent["opinion"] for agent in state["agents"]]
        results["opinions"].append(opinions)
        
        # 记录每个群体的平均观点
        for group_id in range(num_groups):
            group_agents = [agent for agent in state["agents"] if agent["group_id"] == group_id]
            if group_agents:
                avg_opinion = np.mean([agent["opinion"] for agent in group_agents])
                results["group_opinions"][group_id].append(avg_opinion)
                
    return results

def visualize_results(results: dict):
    """可视化模拟结果"""
    # 绘制所有智能体的观点变化
    plt.figure(figsize=(12, 6))
    plt.subplot(1, 2, 1)
    for step_opinions in results["opinions"]:
        plt.scatter(range(len(step_opinions)), step_opinions, alpha=0.1)
    plt.title("所有智能体的观点分布")
    plt.xlabel("智能体ID")
    plt.ylabel("观点值")
    
    # 绘制群体平均观点变化
    plt.subplot(1, 2, 2)
    for group_id, group_opinions in enumerate(results["group_opinions"]):
        plt.plot(group_opinions, label=f"群体 {group_id}")
    plt.title("群体平均观点变化")
    plt.xlabel("时间步")
    plt.ylabel("平均观点值")
    plt.legend()
    
    plt.tight_layout()
    plt.savefig("simulation_results.png")
    plt.close()

if __name__ == "__main__":
    # 模拟参数
    NUM_STEPS = 100
    NUM_AGENTS = 100
    NUM_GROUPS = 3
    NETWORK_DENSITY = 0.1
    
    # 运行模拟
    results = run_simulation(NUM_STEPS, NUM_AGENTS, NUM_GROUPS, NETWORK_DENSITY)
    
    # 保存结果
    with open("simulation_results.json", "w") as f:
        json.dump(results, f)
        
    # 可视化结果
    visualize_results(results) 